// 文件：src/ch08/multi-agent.ts
// import { ChatOpenAI } from "@langchain/openai";
import { ChatOllama } from "@langchain/ollama";
import { PromptTemplate } from "@langchain/core/prompts";
import { Toolkit } from "./ceshi-tool";
import { Readable } from "stream";

interface CollaborateResult {
  individual_results: Record<string, string>;
  final_answer: string;
}
export interface AgentRole {
  name: string;
  description: string;
  expertise: string[];
  tools: string[];
}

export const AGENT_ROLES: AgentRole[] = [
  {
    name: "研究员",
    description: "负责信息收集和事实核查，禁止编造",
    expertise: ["搜索", "验证", "总结"],
    tools: ["search", "get_time"]
  },
  {
    name: "分析师",
    description: "基于真实数据，负责数据分析和洞察，禁止编造数据",
    expertise: ["计算", "统计", "可视化"],
    tools: ["calculator", "analyze_data"]
  },
  {
    name: "协调员",
    description: "整合结果生成结构化总结，禁止添加未验证信息",
    expertise: ["规划", "协调", "总结"],
    tools: ["get_time"]
  }
];

export class MultiAgentSystem {
  private agents: Map<string, { role: AgentRole; llm: ChatOllama; prompt: PromptTemplate }> = new Map();
  private llm: ChatOllama;
  private toolkit: Toolkit; // 工具实例
  constructor() {
    this.llm = new ChatOllama({
      baseUrl: "http://192.168.31.159:11434",
      model: "llama3.1:latest",
    });
    this.toolkit = new Toolkit();
    this.initializeAgents();
  }

  private initializeAgents() {
    for (const role of AGENT_ROLES) {
      this.agents.set(role.name, {
        role,
        llm: this.llm,
        prompt: this.createRolePrompt(role)
      })
    }
  }

  private createRolePrompt(role: AgentRole) {
    return PromptTemplate.fromTemplate(`
你是${role.name}，${role.description}
你的专长：${role.expertise.join(", ")}
可用工具：${role.tools.join(", ")}
工具调用规则：
1. 必须使用可用工具获取数据，禁止编造任何数据（包括数字、品牌、比例等）
2. 工具调用格式严格遵循：
   - 简单参数：工具名("参数内容")，例如 search("2024年新能源汽车销量")
   - JSON参数：工具名({{"key": "value"}})，例如 analyze_data({{"data": "850万辆销量"}})
   注意：JSON格式必须使用双大括号 {{{{ }}}} 包裹
3. 仅能使用指定工具，禁止调用print、get等未注册工具
4. 基于工具返回的真实结果进行分析，不添加无关内容

任务：{task}
当前状态：{state}
其他Agent的意见：{other_opinions}

请输出你的分析（若需调用工具，按格式输出）：
`);
  }

  // 新增：解析智能体输出，判断是否需要调用工具
  private async parseAndExecuteTools(agentOutput: string): Promise<string> {
    // 工具调用格式约定：工具名(参数)
    // 支持格式：
    // - search("关键词")
    // - calculator("10+20")
    // - analyze_data({{"data": "850万辆销量"}}) 或 analyze_data({"data": "850万辆销量"})
    const toolRegex = /(\w+)\(([^)]+)\)/g;
    let matches;
    let result = agentOutput;

    // 提取并执行所有工具调用
    while ((matches = toolRegex.exec(agentOutput)) !== null) {
      const [fullMatch, toolName, rawArgs] = matches;
      const tool = this.toolkit.getTool(toolName);
      if (tool) {
        try {
          // 解析参数：如果是 JSON 格式，尝试解析为对象；否则作为字符串传入
          let parsedArgs: unknown;
          let trimmedArgs = rawArgs.trim();
          
          // 处理双大括号（LLM 可能输出 {{...}} 来避免模板冲突）
          if (trimmedArgs.startsWith('{{') && trimmedArgs.endsWith('}}')) {
            trimmedArgs = trimmedArgs.slice(1, -1); // 移除外层大括号
          }
          
          // 判断是否是 JSON 对象格式
          if (trimmedArgs.startsWith('{') && trimmedArgs.endsWith('}')) {
            try {
              parsedArgs = JSON.parse(trimmedArgs);
            } catch {
              // JSON 解析失败，尝试修复常见问题（如单引号）
              const fixedArgs = trimmedArgs.replace(/'/g, '"');
              try {
                parsedArgs = JSON.parse(fixedArgs);
              } catch {
                // 仍然失败，作为字符串处理
                parsedArgs = trimmedArgs;
              }
            }
          } else {
            // 移除可能的引号
            parsedArgs = trimmedArgs.replace(/^["']|["']$/g, '');
          }
          
          const toolResult = await tool.execute(parsedArgs);
          result = result.replace(fullMatch, `\n[工具调用结果]：${toolResult}\n`);
        } catch (error) {
          result = result.replace(fullMatch, `\n[工具调用失败]：${(error as Error).message}\n`);
        }
      } else {
        result = result.replace(fullMatch, `\n[错误] 未知工具：${toolName}（仅支持${Array.from(this.toolkit.tools.keys()).join(", ")}）\n`);
      }
    }
    return result;
  }

  async collaborate(task: string): Promise<CollaborateResult> {
    const results: Record<string, string> = {};
    const opinions: string[] = [];

    console.log("🔵 === 第一轮：各 Agent 独立分析 ===");
    // 第一轮：各 Agent 独立分析
    // 每个智能体基于自身角色独立分析任务，不参考其他智能体意见，结果存入results。
    for (const [name, agent] of this.agents) {
      console.log(`  → ${name}分析中...`);
      const response = await agent.prompt.pipe(agent.llm).invoke({
        task,
        state: "初始分析阶段",
        other_opinions: "暂无"
      });
      const toolExecutionResult = await this.parseAndExecuteTools(String(response.content));
      results[name] = toolExecutionResult ;
      opinions.push(`${name}: ${toolExecutionResult }`);
    }


    console.log("\n🟡 === 第二轮：协作调整 ===");
    // 第二轮：基于其他 Agent 意见调整
    // 每个智能体参考其他智能体的第一轮意见，调整自己的结论，结果存入results的_final字段。
    for (const [name, agent] of this.agents) {
      console.log(`  → ${name}调整中...`);
      const otherOpinions = opinions.filter(o => !o.startsWith(name));
      const response = await agent.prompt.pipe(agent.llm).invoke({
        task,
        state: "协作调整阶段",
        other_opinions: otherOpinions.join("\n")
      });
      // 再次执行工具调用（如果有）
      const toolResult = await this.parseAndExecuteTools(String(response.content));

      results[`${name}_final`] = toolResult;
    }

    console.log("\n🟢 === 最终整合 ===");
    // 最终整合
    // 由协调员智能体汇总所有智能体的结果，生成最终答案。
    const coordinator = this.agents.get("协调员");
    if (!coordinator) {
      throw new Error("协调员智能体未找到");
    }
    const finalResponse = await coordinator.prompt.pipe(coordinator.llm).invoke({
      task,
      state: "最终整合",
      other_opinions: Object.entries(results).map(([k, v]) => `${k}: ${v}`).join("\n")
    });

    return {
      individual_results: results,
      final_answer: String(finalResponse.content)
    };
  }
}
